import numpy as np
import pyaudio
import pygame
import sys
import audioop
from pygame import gfxdraw

# 初始化pygame
pygame.init()

# 设置显示窗口
WIDTH, HEIGHT = 800, 200
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("立体声声纹显示器")

# 颜色定义
BACKGROUND = (0, 0, 0)  # 黑色背景
LEFT_COLOR = (0, 255, 0)  # 左声道绿色
RIGHT_COLOR = (0, 0, 255)  # 右声道蓝色
CENTER_COLOR = (255, 255, 255)  # 中心线白色

# 音频设置
CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 2
RATE = 44100

# 初始化音频流
p = pyaudio.PyAudio()
stream = p.open(format=FORMAT,
                channels=CHANNELS,
                rate=RATE,
                input=True,
                frames_per_buffer=CHUNK)

# 用于平滑处理的历史数据
history_left = np.zeros(WIDTH)
history_right = np.zeros(WIDTH)

def draw_waveform(left_data, right_data):
    screen.fill(BACKGROUND)
    
    # 绘制中心线
    pygame.draw.line(screen, CENTER_COLOR, (0, HEIGHT//2), (WIDTH, HEIGHT//2), 1)
    
    # 绘制左右声道波形
    for x in range(WIDTH):
        # 左声道 (绿色)
        left_y = int(HEIGHT//2 - left_data[x] * (HEIGHT//2 - 10))
        pygame.gfxdraw.pixel(screen, x, left_y, LEFT_COLOR)
        
        # 右声道 (蓝色)
        right_y = int(HEIGHT//2 - right_data[x] * (HEIGHT//2 - 10))
        pygame.gfxdraw.pixel(screen, x, right_y, RIGHT_COLOR)
        
        # 绘制连接线表示声场方向
        if abs(left_data[x] - right_data[x]) > 0.1:
            mid_y = (left_y + right_y) // 2
            direction = 1 if left_data[x] > right_data[x] else -1
            pygame.draw.line(screen, (255, 255, 0), (x, mid_y), 
                            (x + 5 * direction, mid_y), 1)
    
    # 显示声场方向指示器
    balance = np.mean(left_data) - np.mean(right_data)
    if abs(balance) > 0.05:
        direction = "← 左声道" if balance > 0 else "右声道 →"
        font = pygame.font.SysFont(None, 36)
        text = font.render(direction, True, (255, 255, 255))
        screen.blit(text, (WIDTH//2 - 60, 20))

def process_audio(data):
    # 将原始音频数据转换为numpy数组
    audio_data = np.frombuffer(data, dtype=np.int16)
    
    # 分离左右声道
    left = audio_data[0::2]
    right = audio_data[1::2]
    
    # 计算RMS值作为音量
    left_rms = audioop.rms(left.tobytes(), 2) / 32767
    right_rms = audioop.rms(right.tobytes(), 2) / 32767
    
    # 更新历史数据
    global history_left, history_right
    history_left = np.roll(history_left, -1)
    history_right = np.roll(history_right, -1)
    history_left[-1] = left_rms
    history_right[-1] = right_rms
    
    return history_left, history_right

def main():
    clock = pygame.time.Clock()
    
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
        
        # 读取音频数据
        try:
            data = stream.read(CHUNK, exception_on_overflow=False)
            left, right = process_audio(data)
            draw_waveform(left, right)
        except Exception as e:
            print(f"音频处理错误: {e}")
        
        pygame.display.flip()
        clock.tick(30)
    
    # 清理资源
    stream.stop_stream()
    stream.close()
    p.terminate()
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()